import RPi.GPIO as GPIO
import numpy as np
import time
import datetime as dt
# init func
def init(pin_list):
    GPIO.setmode(GPIO.BCM)
    GPIO.setup(pin_list, GPIO.OUT)
    GPIO.output(pin_list, GPIO.LOW)

# util func
def pulse(pin, duration, is_test=True):
    if duration == 0:
        return
    if is_test:
        print('virtual shake at pin-%s, duration = %f sec' % (pin, duration))
    else:
        print('real shake at pin-%s, duration = %f sec' % (pin, duration))
        step = 0.2
        step_wait = 0.2
        step_num = int(duration / (step + step_wait))
        for i in range(step_num):
            GPIO.output(pin, GPIO.HIGH)
            time.sleep(step)
            GPIO.output(pin, GPIO.LOW)
            time.sleep(step_wait)

def s2date(s):
    return dt.date(int(s[:4]), int(s[4:6]), int(s[7:]))

def is_on_time(start_date_s, stop_date_s):
    fmt = '%Y-%m-%d-%H.%M.%S'
    curr_tm = time.localtime(time.time())
    start_tm = time.strptime(start_date_s, fmt)
    stop_tm = time.strptime(stop_date_s, fmt)
    if curr_tm > start_tm and curr_tm < stop_tm:
        return 1
    else:
        if curr_tm < start_tm:
            diff_sec = time.mktime(start_tm) - time.mktime(curr_tm)
            print('will start after %d sec' % int(diff_sec))
            return 0
        else:
            print('Already out of date.')
            return -1

def update_duration_list(interv_list, stimuli_list, motor_num, timer):
    duration_list = np.zeros(motor_num)
    for i, (interv, stimuli) in enumerate(zip(interv_list, stimuli_list)):
        if interv > 0 and timer % interv == 0: # channel-i, interv/stimuli
            duration_list[i] = stimuli
    return duration_list

def generate_stimuli(motor_num, interv_list_define, stimuli_list_define):
    interv_list = interv_list_define
    if motor_num > len(interv_list):
        interv_list.extend(np.zeros(motor_num - len(interv_list)))
    elif motor_num < len(interv_list):
        interv_list = interv_list[:motor_num]

    stimuli_list = stimuli_list_define
    if motor_num > len(stimuli_list):
        stimuli_list.extend(np.zeros(motor_num - len(stimuli_list)))
    elif motor_num < len(stimuli_list):
        stimuli_list = stimuli_list[:motor_num]
    
    return interv_list, stimuli_list

def batch_pulse(motor_list, duration_list, is_test):
    [pulse(pin_list[motor-1], duration, is_test) for motor, duration in zip(motor_list, duration_list)]

if __name__=='__main__':
    # settings
    is_test = False                      # on test mode, no shake, just output messages
    start_date_s = '2019-12-24-21.00.00' # start date
    stop_date_s  = '2019-12-25-09.00.00' # stop date
    motor_list   = [  6,  4, 10]         # 1-basedi indices
    interv_list  = [ 30, 30, 30]         # interval between stimuli, channel-specific
    stimuli_list = [ 2., 2., 2.]         # stimuli duration, channel-specific

    #stimuli_list = np.array(stimuli_list)
    pin_list = [2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
    motor_num = len(motor_list)
    #interv_list = np.zeros(motor_num) + 60
    interv_list, stimuli_list = generate_stimuli(motor_num, interv_list, stimuli_list)
    duration_list = None
    #duration = 1 # sec of pulse
    step_len = 1.0 # sec to wait for each step
    init(pin_list)
    timer = 1
    pulse_wait = 0.0 # how much time (sec) wait for pulsing
    start_tm = time.time() # initial time point
    stimuli_increment = (6.0 - 2.0) / 720 # gradual increment of stimuli duration
    while(1):
        time_status = is_on_time(start_date_s, stop_date_s)
        if time_status == 1:
            duration_list = update_duration_list(interv_list, stimuli_list, motor_num, timer)
            batch_pulse(motor_list, duration_list, is_test)
            pulse_wait = time.time() - start_tm - timer + 1
            #print(pulse_wait)
            if step_len > pulse_wait:
                time.sleep(step_len - pulse_wait)
            else:
                sec_num = int(pulse_wait)
                debris = pulse_wait - sec_num
                time.sleep(step_len - debris)
                timer += sec_num
            timer += 1
            if timer % 60 == 0:
                stimuli_tmplist = []
                stimuli_tmp = 0.0
                for i, stimuli in enumerate(stimuli_list):
                    stimuli_tmp = 0.0 if stimuli == 0 else stimuli + stimuli_increment * (i*0.5+1)
                    stimuli_tmplist.append(stimuli_tmp)
                stimuli_list = stimuli_tmplist
        elif time_status == 0:
            time.sleep(1)
        else:
            break
